

from flask import Flask, request, jsonify
from flask_sqlalchemy import SQLAlchemy
from flask_mail import Mail, Message
from flask_cors import CORS
from flask_restx import Api, Resource, fields, reqparse
import bcrypt
import random
import string
from datetime import datetime, timedelta
import pymysql
from sqlalchemy import and_  # 导入逻辑运算符

pymysql.install_as_MySQLdb()  # 注册MySQL驱动

# ========== 初始化应用 ==========
app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

# ========== 配置 ==========
# MySQL配置（请替换为你的数据库信息）
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql://root:123456@localhost:3306/flask_auth_restx'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['SQLALCHEMY_ECHO'] = False  # 生产环境关闭SQL日志

# 邮件配置（以QQ邮箱为例，请替换为你的邮箱信息）
app.config['MAIL_SERVER'] = 'smtp.163.com'
app.config['MAIL_PORT'] = 465
app.config['MAIL_USE_SSL'] = True
app.config['MAIL_USERNAME'] = 'qiyuan0529@163.com'  # 发件人邮箱
app.config['MAIL_PASSWORD'] = 'WHdEFKBSirmnvQAJ'   # 邮箱授权码

db = SQLAlchemy(app)
mail = Mail(app)

# ========== 创建API对象 ==========
api = Api(
    app,
    version='1.0',
    title='用户系统 API',
    description='用户注册、登录及邮箱验证接口',
    doc='/docs/',  # 文档路径
    prefix='/api'  # 接口前缀
)

# ========== 数据库模型 ==========
class User(db.Model):
    __tablename__ = 'tb_users'
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    email = db.Column(db.String(100), unique=True, nullable=False)
    password_hash = db.Column(db.String(100), nullable=False)

    def set_password(self, password):
        self.password_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()

    def check_password(self, password):
        return bcrypt.checkpw(password.encode(), self.password_hash.encode())

class VerifyCode(db.Model):
    __tablename__ = 'tb_verify_codes'
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(100), nullable=False)
    code = db.Column(db.String(6), nullable=False)
    type = db.Column(db.SmallInteger, default=1)  # 1=注册，2=找回密码
    expire_time = db.Column(db.DateTime, nullable=False)

# ========== 定义API模型 ==========
# 注册请求模型
register_model = api.model('Register', {
    'username': fields.String(required=True, description='用户名', min_length=3),
    'email': fields.String(required=True, description='邮箱', pattern=r'^[\w\.-]+@[\w\.-]+\.\w+$'),
    'password': fields.String(required=True, description='密码', min_length=8),
    'code': fields.String(required=True, description='6位邮箱验证码')
})

# 登录请求模型
login_model = api.model('Login', {
    'name': fields.String(required=True, description='用户名或邮箱'),
    'password': fields.String(required=True, description='密码')
})

# 通用响应模型
common_response = api.model('CommonResponse', {
    'code': fields.Integer(description='状态码', default=0),
    'message': fields.String(description='状态消息')
})

# 用户详情响应模型
user_detail_model = api.inherit('UserDetail', common_response, {
    'data': fields.Nested(api.model('UserData', {
        'id': fields.Integer(description='用户ID'),
        'username': fields.String(description='用户名'),
        'email': fields.String(description='邮箱')
    }))
})

# ========== 工具函数 ==========
def generate_code():
    return ''.join(random.choices(string.digits, k=6))

def send_email_code(email, code, type=1):
    try:
        msg = Message(
            subject="【用户系统】验证码",
            sender=app.config['MAIL_USERNAME'],
            recipients=[email]
        )
        msg.body = f"您的验证码为：{code}，有效期5分钟。（类型：{'注册' if type==1 else '找回密码'}）"
        mail.send(msg)
        return True
    except Exception as e:
        print(f"邮件发送失败：{str(e)}")
        return False

# ========== 命名空间定义 ==========
user_ns = api.namespace('user', description='用户操作')
email_ns = api.namespace('email', description='邮箱验证')

# ========== 接口实现 ==========
@user_ns.route('/register')
class Register(Resource):
    @api.doc(description='用户注册接口')
    @api.expect(register_model)  # 声明请求模型
    @api.response(201, '注册成功', user_detail_model)  # 声明响应模型
    @api.response(400, '参数错误', common_response)
    @api.response(409, '用户或邮箱已存在', common_response)
    def post(self):
        data = api.payload  # 自动解析请求体

        # 修正：使用 filter + 列对象比较
        verify_code = VerifyCode.query.filter(
            VerifyCode.email == data['email'],
            VerifyCode.code == data['code'],
            VerifyCode.type == 1,
            VerifyCode.expire_time >= datetime.now()
        ).first()

        if not verify_code:
            return {'code': 400, 'message': '验证码错误或已过期'}, 400

        # 检查用户是否存在
        if User.query.filter_by(username=data['username']).first():
            return {'code': 409, 'message': '用户名已存在'}, 409
        if User.query.filter_by(email=data['email']).first():
            return {'code': 409, 'message': '邮箱已注册'}, 409

        # 创建用户
        new_user = User(
            username=data['username'],
            email=data['email']
        )
        new_user.set_password(data['password'])
        db.session.add(new_user)
        db.session.delete(verify_code)  # 验证成功后删除验证码
        db.session.commit()

        return {
            'code': 0,
            'message': '注册成功',
            'data': {
                'id': new_user.id,
                'username': new_user.username,
                'email': new_user.email
            }
        }, 201

@user_ns.route('/login')
class Login(Resource):
    @api.doc(description='用户登录接口')
    @api.expect(login_model)
    @api.response(200, '登录成功', user_detail_model)
    @api.response(401, '认证失败', common_response)
    def post(self):
        data = api.payload
        user = User.query.filter_by(
            username=data['name']
        ).first() or User.query.filter_by(email=data['name']).first()

        if not user or not user.check_password(data['password']):
            return {'code': 401, 'message': '用户名或密码错误'}, 401

        return {
            'code': 0,
            'message': '登录成功',
            'data': {
                'id': user.id,
                'username': user.username,
                'email': user.email
            }
        }, 200

@email_ns.route('/send-code')
class SendCode(Resource):
    @api.doc(
        description='发送邮箱验证码',
        params={'type': '验证码类型（1=注册，2=找回密码，默认1）'}
    )
    @api.param('email', '邮箱地址', required=True)
    @api.response(200, '发送成功', common_response)
    @api.response(400, '参数错误', common_response)
    def get(self):
        parser = reqparse.RequestParser()
        parser.add_argument('email', type=str, required=True, location='args')
        parser.add_argument('type', type=int, default=1, location='args')
        args = parser.parse_args()

        if not args['email']:
            return {'code': 400, 'message': '请提供邮箱地址'}, 400

        # 生成验证码
        code = generate_code()
        expire_time = datetime.now() + timedelta(minutes=5)
        verify_code = VerifyCode(
            email=args['email'],
            code=code,
            type=args['type'],
            expire_time=expire_time
        )
        db.session.add(verify_code)
        db.session.commit()

        # 发送邮件
        if send_email_code(args['email'], code, args['type']):
            return {'code': 0, 'message': '验证码发送成功'}, 200
        else:
            db.session.delete(verify_code)
            db.session.commit()
            return {'code': 400, 'message': '验证码发送失败'}, 400

# ========== 初始化数据库 ==========
with app.app_context():
    db.create_all()  # 创建表结构

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8888, debug=True)